import pathlib
import typing

import matplotlib.pyplot
import torch

import models
import shared.checkpoint


def plot_checkpoint(checkpoint: pathlib.Path,
                    times: typing.Sequence[float] = (0.59, 0.79, 0.98)) \
        -> matplotlib.pyplot.Figure:
    states = torch.load(checkpoint)

    model = models.get_model()
    model.load_state_dict(states["model"])
    model.eval()

    figure = matplotlib.pyplot.figure(figsize=(8, 5), dpi=110)

    def _plot_contour():
        axes = figure.add_subplot(2, 1, 1)

        x = torch.linspace(-5, 5, 1000)
        t = torch.linspace(0, torch.pi / 2, 1000)

        x, t = torch.meshgrid(x, t, indexing="xy")
        xt = torch.stack([x, t], dim=-1)

        h = model(xt)
        h = (h[:, :, 0] ** 2 + h[:, :, 1] ** 2) ** (1 / 2)

        contour = axes.contourf(torch.detach(t), torch.detach(x), torch.detach(h),
                                levels=80, cmap=matplotlib.pyplot.get_cmap("YlGnBu"))

        axes.set_title("$|h(t, x)|$")
        axes.set_xlabel("$t$")
        axes.set_ylabel("$x$")
        figure.colorbar(contour, fraction=0.08, pad=0.01)

    _plot_contour()

    def _plot_times():
        for i, time in enumerate(times, start=len(times) + 1):
            axes = figure.add_subplot(2, len(times), i)

            x = torch.linspace(-5, 5, 1000)
            t = torch.ones_like(x) * time
            xt = torch.stack([x, t], dim=-1)

            h = model(xt)
            h = (h[:, 0] ** 2 + h[:, 1] ** 2) ** (1 / 2)

            axes.set_xlim(left=-5, right=5)
            axes.set_ylim(bottom=0, top=5)
            axes.plot(torch.detach(x), torch.detach(h))

            axes.set_title(f"$t={time}$")
            axes.set_xlabel("$x$")
            axes.set_ylabel("$|h(t, x)|$")

    _plot_times()

    figure.subplots_adjust(wspace=0.357, hspace=0.502)
    return figure


def main():
    checkpoint_manager = shared.checkpoint.CheckpointManager(pathlib.Path("./checkpoints"))
    checkpoint = checkpoint_manager.last_existed()
    assert checkpoint is not None

    figure = plot_checkpoint(checkpoint[1])
    figure.show()
    matplotlib.pyplot.pause(0)


if __name__ == "__main__":
    main()
